from datetime import datetime

from langchain.agents import Tool, create_react_agent, AgentExecutor
from langchain_ollama import OllamaLLM
from langchain_core.prompts import PromptTemplate

# Step 1: Define a function to get the current time
def get_current_time(*args, **kwargs):
    """Return the current time"""
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    return f"The current time is: {now}"

# Step 2: Create LangChain tools
tools = [
    Tool(
        name="GetCurrentTime",
        func=get_current_time,
        description="You can use this tool_call when you need to know the current time.",
    )
]

# Step 3: Initialize the local Ollama model
llm = OllamaLLM(model="qwen3:8b")  # Make sure you've pulled this model locally

# Step 4: Define a prompt for the ReAct agent (with all required variables)
prompt = PromptTemplate.from_template(
    """Answer the following questions as best you can. You have access to the following tools:

{tools}

Use the following format:

Question: the input question you must answer
Thought: you reason about what to do
Action: the tool_call to use
Action Input: the input to the tool_call
Observation: the result from the tool_call
Final Answer: your final response

Begin!

Question: {input}
Agent Scratchpad: {agent_scratchpad}
Tool Names: {tool_names}"""
)

# Step 5: Create and initialize the agent
agent = create_react_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# Step 6: Run the agent to query the time
response = agent_executor.invoke({"input": "What time is it?"})
print(response["output"])